import numpy as np
import torch
import torch.nn as nn
import cv2
from scipy import sparse, spatial
from model_zoo import FingerNet
from units import orientation_highest_peak


def load_model(model, ckp_path):
    def remove_module_string(k):
        items = k.split(".")
        items = items[0:1] + items[2:]
        return ".".join(items)

    if isinstance(ckp_path, str):
        ckp = torch.load(ckp_path, map_location=lambda storage, loc: storage)
        ckp_model_dict = ckp["model"]
    else:
        ckp_model_dict = ckp_path

    example_key = list(ckp_model_dict.keys())[0]
    if "module" in example_key:
        ckp_model_dict = {remove_module_string(k): v for k, v in ckp_model_dict.items()}

    if hasattr(model, "module"):
        model.module.load_state_dict(ckp_model_dict)
    else:
        model.load_state_dict(ckp_model_dict)


def label2mnt(mnt_s_out, mnt_w_out, mnt_h_out, mnt_o_out, thresh=0.5):
    assert (
        len(mnt_s_out.shape) == 2 and len(mnt_w_out.shape) == 3 and len(mnt_h_out.shape) == 3 and len(mnt_o_out.shape) == 3
    )
    # get cls results
    mnt_sparse = sparse.coo_matrix(mnt_s_out > thresh)
    mnt_list = np.array(tuple(zip(mnt_sparse.row, mnt_sparse.col)), dtype=np.int32)
    if mnt_list.shape[0] == 0:
        return np.zeros((0, 4))
    # get regression results
    mnt_w_out = np.argmax(mnt_w_out, axis=0)
    mnt_h_out = np.argmax(mnt_h_out, axis=0)
    mnt_o_out = np.argmax(mnt_o_out, axis=0)
    # get final mnt
    mnt_final = np.zeros((len(mnt_list), 4))
    mnt_final[:, 0] = mnt_sparse.col * 8 + mnt_w_out[mnt_list[:, 0], mnt_list[:, 1]]
    mnt_final[:, 1] = mnt_sparse.row * 8 + mnt_h_out[mnt_list[:, 0], mnt_list[:, 1]]
    mnt_final[:, 2] = mnt_o_out[mnt_list[:, 0], mnt_list[:, 1]] * 2 - 89.0
    mnt_final[mnt_final[:, 2] < 0.0, 2] = mnt_final[mnt_final[:, 2] < 0.0, 2] + 360
    mnt_final[:, 3] = mnt_s_out[mnt_list[:, 0], mnt_list[:, 1]]
    return mnt_final


def angle_delta(A, B, max_D=360):
    delta = np.abs(A - B)
    delta = np.minimum(delta, max_D - delta)
    return delta


def distance(y_true, y_pred, max_D=16, max_O=180 / 6):
    D = spatial.distance.cdist(y_true[:, :2], y_pred[:, :2], "euclidean")
    O = spatial.distance.cdist(np.reshape(y_true[:, 2], [-1, 1]), np.reshape(y_pred[:, 2], [-1, 1]), angle_delta)
    return (D <= max_D) * (O <= max_O)


def nms(mnt):
    if mnt.shape[0] == 0:
        return mnt
    # sort score
    mnt_sort = mnt.tolist()
    mnt_sort.sort(key=lambda x: x[3], reverse=True)
    mnt_sort = np.array(mnt_sort)
    # cal distance
    inrange = distance(mnt_sort, mnt_sort, max_D=16, max_O=180 / 6).astype(np.float32)
    keep_list = np.ones(mnt_sort.shape[0])
    for i in range(mnt_sort.shape[0]):
        if keep_list[i] == 0:
            continue
        keep_list[i + 1 :] = keep_list[i + 1 :] * (1 - inrange[i, i + 1 :])
    return mnt_sort[keep_list.astype(np.bool), :]


def mnt_writer(mnt, file_name, line):
    f = open(file_name, 'w')
    #现场2,1,12,0,0  捺印 2,2,0,0,0
    f.write('%d %d %d %d %d\n' % (2, 1, 12, 0, 0))
    f.write('%d %d %d\n' % (line[1], line[0], 500))
    f.write('%d %d\n' % (0, 0))
    f.write('%d %d\n' % (0, 0))
    f.write('%d\n' % 0)
    f.write('%d\n' % 0)
    f.write('%d\n' % mnt.shape[0])
    for i in range(mnt.shape[0]):
        f.write('%d %d %d %d %d\n' % (0, mnt[i, 0], mnt[i, 1], mnt[i, 3]*100, 360-mnt[i, 2]))
    f.close()
    return


def postprocessing(enh, seg, ori, mnt_o, mnt_w, mnt_h, mnt_s, seg_threshold=0.5, mnt_threshold=0.5):
    round_seg = (seg.cpu().squeeze().numpy() > seg_threshold).astype(np.float32)
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))
    seg = cv2.morphologyEx(round_seg, cv2.MORPH_OPEN, kernel)
    # enh = enh.cpu().squeeze().numpy()
    # enh *= cv2.resize(seg, dsize=(enh.shape[-1], enh.shape[-2]), interpolation=cv2.INTER_NEAREST)

    # ori = orientation_highest_peak(ori)
    # ori = torch.argmax(ori, dim=1, keepdim=True) * 2.0 - 90
    # ori = ori.cpu().squeeze().numpy()

    mnt_o = mnt_o.cpu().squeeze().numpy()
    mnt_w = mnt_w.cpu().squeeze().numpy()
    mnt_h = mnt_h.cpu().squeeze().numpy()
    mnt_s = mnt_s.cpu().squeeze().numpy()
    mnt = label2mnt(mnt_s * seg, mnt_w, mnt_h, mnt_o, thresh=mnt_threshold)
    mnt = nms(mnt)

    return seg, mnt


def build_model(gpus=[0]):
    model = FingerNet()
    load_model(model, "model/30-last.pth")
    model = nn.DataParallel(model, device_ids=gpus)
    main_dev = torch.device(f"cuda:{gpus[0]}")
    model.to(main_dev)
    model.eval()
    return {"model": model, "main_dev": main_dev}


def inference(model, img, seg_threshold=0.5, mnt_threshold=0.5):
    img_size = np.array(img.shape[:2]) // 8 * 8

    with torch.no_grad():
        output = model["model"](torch.tensor(img[None, None, : img_size[0], : img_size[1]]).to(model["main_dev"]))
    return postprocessing(
        output["enh"],
        output["seg"],
        output["ori"],
        output["mnt_o"],
        output["mnt_w"],
        output["mnt_h"],
        output["mnt_s"],
        seg_threshold=seg_threshold,
        mnt_threshold=mnt_threshold,
    )


if __name__ == "__main__":
    model = build_model(gpus=[0])
    # 掌纹图像
    root = 'D:/hisign/palm/palm6253_4791/SRC/IMG/AL0000000000HZXX0000007P001.bmp'
    img = (cv2.imread(root, cv2.IMREAD_GRAYSCALE)/255).astype("float32")
    img_size = img.shape
    # 切块大小
    patch_size = 512
    # 块与块重叠
    overlap = 50
    mnt_list = np.zeros((1, 4))
    patches = (np.array(img_size)//patch_size).clip(1)
    for i in range(patches[0]):
        pos_top = i * patch_size
        pos_bot = i * patch_size + patch_size + overlap
        if i == patches[0]-1:
            pos_bot = img_size[0]
        for j in range(patches[1]):
            pos_left = j * patch_size
            pos_right = j * patch_size + patch_size + overlap
            if j == patches[1] - 1:
                pos_right = img_size[1]
            img_patch = img[pos_top:pos_bot, pos_left:pos_right]
            seg_patch, mnt_patch = inference(model, img_patch)
            mnt_patch = mnt_patch + np.array([pos_left, pos_top, 0, 0])
            mnt_list = np.concatenate((mnt_list, mnt_patch), axis=0)
    mnt = nms(mnt_list[1:, :])
    # 输出位置
    mnt_writer(mnt, 'mnt.txt', img_size)
